/*
 * Parse CMS Enveloped data
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include "cms.h"
#include <string.h>
#include <stdlib.h>
#include <openssl/asn1.h>
#include <openssl/crypto.h>
#include <openssl/pem.h>
#include <openssl/cms.h>
#include <openssl/err.h>
#include <stdbool.h>

#define ASN1_ERROR                      0x80

#define ASN1_CONS_SEQUENCE(rc, tag) \
    (!((rc) & ASN1_ERROR) && ((rc) & V_ASN1_CONSTRUCTED) && ((tag) == V_ASN1_SEQUENCE))

#define ASN1_PRIM_OBJECT(rc, tag) \
    (!((rc) & ASN1_ERROR) && !((rc) & V_ASN1_CONSTRUCTED) && ((tag) == V_ASN1_OBJECT))

#define ASN1_PRIM_INTEGER(rc, tag) \
    (!((rc) & ASN1_ERROR) && !((rc) & V_ASN1_CONSTRUCTED) && ((tag) == V_ASN1_INTEGER))

#define ASN1_CONS_CONT(rc, xclass) \
    (!((rc) & ASN1_ERROR) && ((rc) & V_ASN1_CONSTRUCTED) && \
    (((xclass) & V_ASN1_CONTEXT_SPECIFIC) == V_ASN1_CONTEXT_SPECIFIC))

#define ASN1_PRIM_CONT(rc, xclass) \
    (!((rc) & ASN1_ERROR) && !((rc) & V_ASN1_CONSTRUCTED) && \
    (((xclass) & V_ASN1_CONTEXT_SPECIFIC) == V_ASN1_CONTEXT_SPECIFIC))

#define ASN1_PRIM_OCTET_STRING(rc, tag) \
    (!((rc) & ASN1_ERROR) && ((tag) == V_ASN1_OCTET_STRING))

#define ASN1_CONS_SET(rc, tag) \
    (!((rc) & ASN1_ERROR) && ((rc) & V_ASN1_CONSTRUCTED) && ((tag) == V_ASN1_SET))

static bool stream_buf_is_valid(const struct stream_buf *binary_in)
{
    if (binary_in == NULL ||
       !((binary_in->capacity == 0 && binary_in->len == 0 && binary_in->buffer == NULL) ||
        (binary_in->capacity > 0 && binary_in->len <= binary_in->capacity &&
        binary_in->buffer != NULL)))
        return false;

    return true;
}

static int check_enveloped_data_type(const u_char **cur_p, long *length,
                                     ASN1_OBJECT **obj)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_RESP_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RESP_ERR;

    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OBJECT((unsigned int)rc, tag))
        return CMS_RESP_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RESP_ERR;

    *obj = d2i_ASN1_OBJECT(obj, &ori_p, data_len + (*cur_p - ori_p));
    if (*obj == NULL || OBJ_obj2nid(*obj) != NID_pkcs7_enveloped)
        return CMS_RESP_ERR;

    if ((*length -= data_len) < 0)
        return CMS_RESP_ERR;

    *cur_p += data_len;

    return CMS_SUCCESS;
}

static int check_enveloped_data_version(const u_char **cur_p, long *length,
                                        ASN1_INTEGER **ver)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* cons: cont [ 0 ] */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_CONT((unsigned int)rc, (unsigned int)xclass) || tag != 0)
        return CMS_ENVELOPE_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_ENVELOPE_VER_ERR;

    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_ENVELOPE_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_ENVELOPE_VER_ERR;

    /* prim: integer */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_INTEGER((unsigned int)rc, tag))
        return CMS_ENVELOPE_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_ENVELOPE_VER_ERR;

    *ver = d2i_ASN1_INTEGER(ver, &ori_p, data_len + (*cur_p - ori_p));
    if (*ver == NULL || ASN1_INTEGER_get(*ver) != ENVELOPED_VERSION)
        return CMS_ENVELOPE_VER_ERR;

    if ((*length -= data_len) < 0)
        return CMS_ENVELOPE_VER_ERR;

    *cur_p += data_len;

    return CMS_SUCCESS;
}

static int get_recipient_key_version(const u_char **cur_p, long *length,
                                     ASN1_INTEGER **ver)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* cons: set */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SET((unsigned int)rc, tag))
        return CMS_RECIP_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RECIP_VER_ERR;

    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_RECIP_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RECIP_VER_ERR;

    /* KeyTransRecipientInfo: version */
    /* Only subjectKeyIdentifier is supported */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_INTEGER((unsigned int)rc, tag))
        return CMS_RECIP_VER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RECIP_VER_ERR;

    *ver = d2i_ASN1_INTEGER(ver, &ori_p, data_len + (*cur_p - ori_p));
    if (*ver == NULL || ASN1_INTEGER_get(*ver) != ENVELOPED_RECIPIENT_VERSION)
        return CMS_RECIP_VER_ERR;

    if ((*length -= data_len) < 0)
        return CMS_RECIP_VER_ERR;

    *cur_p += data_len;

    return CMS_SUCCESS;
}

static int get_recipient_rid_algo(const u_char **cur_p, long *length,
                                  ASN1_OBJECT **obj)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* KeyTransRecipientInfo: rid & algo */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_CONT((unsigned int)rc, (unsigned int)xclass))
        return CMS_RECIP_ALG_ERR;

    if ((*length -= (*cur_p - ori_p + data_len)) < 0)
        return CMS_RECIP_ALG_ERR;

    *cur_p += data_len;

    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_RECIP_ALG_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RECIP_ALG_ERR;

    /* Asymmetric ALGO. RSA-OAEP in this case */
    /* prim: object */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OBJECT((unsigned int)rc, tag))
        return CMS_RECIP_ALG_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_RECIP_ALG_ERR;

    *obj = d2i_ASN1_OBJECT(obj, &ori_p, data_len + (*cur_p - ori_p));
    if (*obj == NULL || OBJ_obj2nid(*obj) != NID_rsaesOaep)
        return CMS_RECIP_ALG_ERR;

    if ((*length -= data_len) < 0)
        return CMS_RECIP_ALG_ERR;

    *cur_p += data_len;

    return CMS_SUCCESS;
}

static int get_recipient_encrypted_key(const u_char **cur_p, long *length,
                                       ASN1_OCTET_STRING **symmkey,
                                       struct stream_buf *cipher_symm_key)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* recipient encrypted key */
    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_SYMMKEY_ERR;

    if ((*length -= (*cur_p - ori_p + data_len)) < 0)
        return CMS_SYMMKEY_ERR;

    *cur_p += data_len;

    /* prim: octet string */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OCTET_STRING((unsigned int)rc, tag))
        return CMS_SYMMKEY_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_SYMMKEY_ERR;

    *symmkey = d2i_ASN1_OCTET_STRING(NULL, &ori_p, data_len + (*cur_p - ori_p));
    if ((*symmkey) == NULL || (*symmkey)->length <= 0)
        return CMS_SYMMKEY_ERR;

    if ((*length -= (*symmkey)->length) < 0)
        return CMS_SYMMKEY_ERR;

    *cur_p += (*symmkey)->length;

    cipher_symm_key->capacity = cipher_symm_key->len = (*symmkey)->length;
    cipher_symm_key->buffer = malloc(cipher_symm_key->capacity);
    if (!(cipher_symm_key->buffer))
        return CMS_INTERNAL_ERR;
    memcpy(cipher_symm_key->buffer, (*symmkey)->data, (*symmkey)->length);

    return CMS_SUCCESS;
}

static int get_encrypted_key(const u_char **cur_p, long *length,
                             ASN1_OBJECT **obj, ASN1_INTEGER **ver,
                             ASN1_OCTET_STRING **symmkey,
                             struct stream_buf *cipher_symm_key)
{
    int rc;

    rc = get_recipient_key_version(cur_p, length, ver);
    if (rc != CMS_SUCCESS)
        return rc;

    rc = get_recipient_rid_algo(cur_p, length, obj);
    if (rc != CMS_SUCCESS)
        return rc;

    return get_recipient_encrypted_key(cur_p, length, symmkey, cipher_symm_key);
}

static int check_encrypted_content_type(const u_char **cur_p, long *length,
                                        ASN1_OBJECT **obj)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_CONTENT_TYPE_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_TYPE_ERR;

    /* prim: object */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OBJECT((unsigned int)rc, tag))
        return CMS_CONTENT_TYPE_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_TYPE_ERR;

    *obj = d2i_ASN1_OBJECT(obj, &ori_p, data_len + (*cur_p - ori_p));
    if (*obj == NULL || OBJ_obj2nid(*obj) != NID_pkcs7_data)
        return CMS_CONTENT_TYPE_ERR;

    if ((*length -= data_len) < 0)
        return CMS_CONTENT_TYPE_ERR;

    *cur_p += data_len;

    return CMS_SUCCESS;
}

static int get_iv(const u_char **cur_p, long *length, ASN1_OBJECT **obj,
    ASN1_OCTET_STRING **key_iv, struct stream_buf *symm_key_iv)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;

    /* Validate that we have AES256-CBC in the Content */
    /* cons: sequence */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_SEQUENCE((unsigned int)rc, tag))
        return CMS_CONTENT_ALGO_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_ALGO_ERR;

    /* prim: object */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OBJECT((unsigned int)rc, tag))
        return CMS_CONTENT_ALGO_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_ALGO_ERR;

    *obj = d2i_ASN1_OBJECT(obj, &ori_p, data_len + (*cur_p - ori_p));
    if (*obj == NULL || OBJ_obj2nid(*obj) != NID_aes_256_cbc)
        return CMS_CONTENT_ALGO_ERR;

    if ((*length -= data_len) < 0)
        return CMS_CONTENT_ALGO_ERR;

    *cur_p += data_len;

    /* prim: octet string */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_PRIM_OCTET_STRING((unsigned int)rc, tag))
        return CMS_CONTENT_IV_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_ALGO_ERR;

    *key_iv = d2i_ASN1_OCTET_STRING(NULL, &ori_p, data_len + (*cur_p - ori_p));
    if (*key_iv == NULL || (*key_iv)->length <= 0)
        return CMS_CONTENT_IV_ERR;

    if ((*length -= (*key_iv)->length) < 0)
        return CMS_CONTENT_ALGO_ERR;

    *cur_p += (*key_iv)->length;

    symm_key_iv->capacity = symm_key_iv->len = (*key_iv)->length;
    symm_key_iv->buffer = malloc(symm_key_iv->capacity);
    if (!(symm_key_iv->buffer))
        return CMS_INTERNAL_ERR;
    memcpy(symm_key_iv->buffer, (*key_iv)->data,
        (*key_iv)->length);

    return CMS_SUCCESS;
}

static int get_encrypted_content(const u_char **cur_p, long *length,
                                 ASN1_OCTET_STRING **text,
                                 struct stream_buf *cipher_text)
{
    const u_char *ori_p = NULL;
    long data_len;
    int tag;
    int xclass;
    int rc;
    u_char *cur_cipher_text_pos = NULL;

    /* cons: cont [ 0 ] */
    ori_p = *cur_p;
    rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
    if (!ASN1_CONS_CONT((unsigned int)rc, (unsigned int)xclass) || tag != 0)
        return CMS_CONTENT_CIPHER_ERR;

    if ((*length -= (*cur_p - ori_p)) < 0)
        return CMS_CONTENT_CIPHER_ERR;

    cipher_text->capacity = *length;
    cipher_text->buffer = cur_cipher_text_pos = malloc(cipher_text->capacity);
    if (!(cipher_text->buffer))
        return CMS_INTERNAL_ERR;

    while (*length > 0) {
        /* prim: OCTET STRING or prim EOC */
        ori_p = *cur_p;
        rc = ASN1_get_object(cur_p, &data_len, &tag, &xclass, *length);
        if ((tag == V_ASN1_EOC) && (xclass == 0))
            break;

        if (!ASN1_PRIM_OCTET_STRING((unsigned int)rc, tag))
            return CMS_CONTENT_CIPHER_ERR;

        if ((*length -= (*cur_p - ori_p)) < 0)
            return CMS_CONTENT_CIPHER_ERR;

        *text = d2i_ASN1_OCTET_STRING(NULL, &ori_p,
            data_len + (*cur_p - ori_p));
        if (*text == NULL || (*text)->length <= 0)
            return CMS_CONTENT_CIPHER_ERR;

        if ((*length -= (*text)->length) < 0)
            return CMS_CONTENT_CIPHER_ERR;

        *cur_p += (*text)->length;
        cipher_text->len += (*text)->length;

        memcpy(cur_cipher_text_pos, (*text)->data, (*text)->length);
        cur_cipher_text_pos += (*text)->length;
    }

    return CMS_SUCCESS;
}

int parse_response_enveloped_data(const struct stream_buf *ciphertext_for_recipient,
    struct stream_buf *cipher_symm_key, struct stream_buf *symm_key_iv,
    struct stream_buf *cipher_text)
{
    const u_char *cur_p;
    long length;
    ASN1_OBJECT *obj = NULL;
    ASN1_INTEGER *ver = NULL;
    ASN1_OCTET_STRING *symmkey = NULL;
    ASN1_OCTET_STRING *key_iv = NULL;
    ASN1_OCTET_STRING *text = NULL;
    int rc;

    if (!stream_buf_is_valid(ciphertext_for_recipient))
        return CMS_INVALID_ARGUMENT;

    cur_p = ciphertext_for_recipient->buffer;
    length = ciphertext_for_recipient->len;

    /* Validate that this is PKCS7 Enveloped Data type we support */
    rc = check_enveloped_data_type(&cur_p, &length, &obj);
    if (rc != CMS_SUCCESS)
        goto err;

    /* PKCS7 Enveloped data. version */
    /* Validate that this is Enveloped Data version we support */
    rc = check_enveloped_data_version(&cur_p, &length, &ver);
    if (rc != CMS_SUCCESS)
        goto err;

    /* Todo: Originator Info. Optional, but if present, consume it */

    /* PKCS7 Enveloped data. recipientInfos */
    rc = get_encrypted_key(&cur_p, &length, &obj, &ver,
        &symmkey, cipher_symm_key);
    if (rc != CMS_SUCCESS)
        goto err;

    /* EncryptedContentInfo */
    /* Validate that this is PKCS7 Data type */
    rc = check_encrypted_content_type(&cur_p, &length, &obj);
    if (rc != CMS_SUCCESS)
        goto err;

    /* Fetch the IV */
    rc = get_iv(&cur_p, &length, &obj, &key_iv, symm_key_iv);
    if (rc != CMS_SUCCESS)
        goto err;

    /* Fetch the encrypted content */
    rc = get_encrypted_content(&cur_p, &length, &text, cipher_text);

err:
    ASN1_OBJECT_free(obj);
    ASN1_INTEGER_free(ver);
    ASN1_OCTET_STRING_free(symmkey);
    ASN1_OCTET_STRING_free(key_iv);
    ASN1_OCTET_STRING_free(text);
    return rc;
}

int response_cipher_decrypt(struct stream_buf *cipher_text,
                            struct stream_buf *plain_symm_key,
                            struct stream_buf *symm_key_iv,
                            struct stream_buf *plain_text)
{
    EVP_CIPHER_CTX *ctx = NULL;
    int normal_len;
    int final_len;
    int rc = CMS_SUCCESS;

    if (!stream_buf_is_valid(cipher_text) ||
        !stream_buf_is_valid(plain_symm_key) ||
        !stream_buf_is_valid(symm_key_iv))
        return CMS_INVALID_ARGUMENT;

    if (plain_symm_key->len != (u_int)EVP_CIPHER_key_length(EVP_aes_256_cbc()))
        return CMS_SYMMKEY_ERR;

    ctx = EVP_CIPHER_CTX_new();
    if (ctx == NULL)
        return CMS_DECRYPT_ERR;

    /* Setup the decryption context */
    if (!EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, plain_symm_key->buffer,
        symm_key_iv->buffer)) {
        EVP_CIPHER_CTX_free(ctx);
        return CMS_DECRYPT_ERR;
    }

    /* out_text_len = in_len + cipher block size - 1 */
    uint8_t out_text[cipher_text->len + EVP_CIPHER_CTX_block_size(ctx)];
    if (!EVP_DecryptUpdate(ctx, out_text, &normal_len,
        cipher_text->buffer, cipher_text->len) ||
        !EVP_DecryptFinal_ex(ctx, &out_text[normal_len], &final_len)) {
        ERR_print_errors_fp(stderr);
        rc = CMS_DECRYPT_ERR;
        goto err;
    }

    /* plain_text */
    if (plain_text->capacity < (u_int)(normal_len + final_len)) {
        rc = CMS_BUFF_SMALL;
        goto err;
    }

    plain_text->len = (u_int)(normal_len + final_len);
    memcpy(plain_text->buffer, out_text, normal_len + final_len);

err:
    EVP_CIPHER_CTX_free(ctx);

    return rc;
}
